#!/usr/bin/env python

import sys
from OpenGL.GL import *
from OpenGL.GLE import *
from OpenGL.GLU import *
from OpenGL.GLUT import *
import pyui

from threading import Timer

from math  import log
from numpy import array
from numpy import concatenate
from numpy import ones
from numpy import reshape
from numpy import zeros

from scikits.ann import kdtree

import psyco
psyco.full()

class OpenGLDataView(object):
  """Objet de visualisation OpenGL"""
  
  def __init__(self, data, knn=0, neighbours=None, colors=None, width=1280, height=800, fullscreen=True):
    """Initialization"""
    
    self.width      = width
    self.height     = height
    self.fullscreen = fullscreen
    
    self.initView()
    self.initData(data)
    self.initNeighbours(knn, neighbours)
    self.initUI()
    self.initGL()
    self.initGLData()
    
    if colors != None:
      self.initGLColors(colors)
    
    self.setStatus('Initialisation done')
  
  def initView(self):
    """
    Set up initial view (rotation, translation, zoom)
    """
    # Set up initial rotation,
    self.rotx,   self.roty   = 0, 0
    # translation,
    self.transx, self.transy = 0, 0
    # zoom,
    self.zoom     = 1
  
  def initData(self, data):
    # Set up data points
    self.n   = data.shape[0]
    self.dim = data.shape[1]
    
    if self.dim < 3:
      data = concatenate((data, zeros((self.n, 3 - self.dim))), 1)
      self.dim = 3
    
    self.data = data
    
    # Calculate bounding hypersphere for the data
    currentMin = float('infinity')
    currentMax = float('-infinity')
    
    for ligne in self.data:
      newMin = min(ligne)
      newMax = max(ligne)
      
      if newMin < currentMin:
        currentMin = newMin
      else:
        if newMax > currentMax:
          currentMax = newMax
    
    # Handle reshaping properly and initialize projection matrix
    self.bounds = 2 * max((abs(x) for x in (currentMin, currentMax)))
  
  def initNeighbours(self, knn, neighbours):
    """Initialise neighbours"""
    if type(knn) != int or knn < 0:
      raise ValueError, 'knn parameter should be a positive integer'
    
    if knn > 0:
      if neighbours != None:
        raise ValueError, 'knn and neighbours parameters can not be set simultaneously'
      
      neighbours = kdtree(self.data).knn(self.data, knn)[0]
    
    if neighbours == None:
      self.k = 0
    else:
      if neighbours.shape[0] != self.n:
        error = "neighbours parameter should be a (%d, k) array, instead received a %s array."%(self.n, neighbours.shape)
        raise IndexError, error
      
      self.k = neighbours.shape[1]
    
    if self.k > 0:
      # Add neighbours to graph
      self.data = concatenate((
        self.data,
        reshape(array([
          (self.data[i], self.data[neighbours[i][j]])
            for j in range(0, self.k) for i in range(0, self.n)]),
          (self.n * self.k * 2, self.dim))))
      
      self.displayNeighbours = True
  
  def initDimFrame(self):
    x, y = 870, 0
    
    self.dimFrame = pyui.widgets.Frame(x, y, 400, 100, "Select dimensions to be displayed")
    self.dimFrame.setLayout(pyui.layouts.TableLayoutManager(4, 5, padding=0))
    
    # Dimension selectors
    self.dimEditSource = [None, None, None]
    self.dimEditTarget = [None, None, None]
    
    for i in range(0, 3):
      self.dimEditSource[i] = pyui.widgets.NumberEdit(str(i),                self.dim, self.updateDisplayedData, False)
      self.dimEditTarget[i] = pyui.widgets.NumberEdit(str(self.dim - i - 1), self.dim, self.updateDisplayedData, False)
      
      self.dimFrame.addChild(self.dimEditSource[i], (1, i, 1, 1))
      self.dimFrame.addChild(self.dimEditTarget[i], (2, i, 1, 1))
    
    self.dimFrame.addChild(pyui.widgets.Label('Source ( 0 to %d )'%(self.dim - 1)), (0, 1, 1, 1))
    self.dimFrame.addChild(pyui.widgets.Label('Target ( 0 to %d )'%(self.dim - 1)), (3, 1, 1, 1))
    
    # Progress slider
    self.dimFrame.addChild(pyui.widgets.Label('Progress'), (0, 3, 1, 1))
    self.progressSlider = pyui.widgets.SliderBar(self.updateDisplayedData, 100, 0)
    self.dimFrame.addChild(self.progressSlider, (1, 3, 2, 1))
    self.dimFrame.addChild(pyui.widgets.Button('Switch', handler=self.doSwitchDimensions, fgColor=(0,0,0,255)), (3, 2, 1, 1))
    self.dimFrame.addChild(pyui.widgets.Button('Reset', handler=self.doResetProgress, fgColor=(0,0,0,255)), (3, 3, 1, 1))
    
    # Speed slider
    self.dimFrame.addChild(pyui.widgets.Label('Speed'),                 (0, 4, 1, 1))
    self.speedSlider = pyui.widgets.SliderBar(None, 20, 1)
    self.dimFrame.addChild(self.speedSlider, (1, 4, 2, 1))
    
    # Play button
    self.animating = False
    self.playButton = pyui.widgets.Button('Play', handler=self.doPlay, fgColor=(0,0,0,255))
    self.dimFrame.addChild(self.playButton, (3, 4, 1, 1))
    
    def doHideDimFrame():
      self.dimFrame.setShow(False)
      self.dimFrame.moveto(x, y)
    
    self.dimFrame._pyuiCloseButton = doHideDimFrame
    
    self.dimFrame.pack()
  
  def initHelpFrame(self):
    numHelpItems = len(self.keyHandler)
    
    height = 17 * numHelpItems
    
    maxTextWidth = 0
    
    for key in self.keyHandler:
      maxTextWidth = max(maxTextWidth, len(self.keyHandler[key][1]))
    
    width  = 6 * maxTextWidth + 20
    
    x, y = self.width - 10 - width, self.height - 30 - height
    
    self.helpFrame = pyui.widgets.Frame(x, y, width, height, "Keyboard Help")
    self.helpFrame.setLayout(pyui.layouts.GridLayoutManager(1, numHelpItems, padding=1))
    
    for key in self.keyHandler:
      self.helpFrame.addChild(pyui.widgets.Label(
        '%s : %s'%(key, self.keyHandler[key][1]),
        font="courier"))
    
    def doHideHelpFrame():
      self.helpFrame.setShow(False)
      self.helpFrame.moveto(x, y)
      self.helpVisible = False
    
    self.helpFrame._pyuiCloseButton = doHideHelpFrame
    
    self.helpFrame.pack()
    
    doHideHelpFrame()
  
  def initActionFrame(self):
    self.actionFrame = pyui.base.Window(0, 750, 700, 30, "Actions")
    self.actionFrame.setLayout(pyui.layouts.FlowLayoutManager(1))
    
    self.actionFrame.addChild(pyui.widgets.Button('Quit', handler=self.doQuit, fgColor=(0,0,0,255)))
    self.actionFrame.addChild(pyui.widgets.Button('Save image', handler=self.doSave, fgColor=(0,0,0,255)))
    self.actionFrame.addChild(pyui.widgets.Label("   Zoom (press 'z' for mouse control):"))
    self.zoomSlider = pyui.widgets.SliderBar(self.changeZoom, 200, 25)
    self.zoomSlider.resize(200, 20)
    self.actionFrame.addChild(self.zoomSlider)
    self.statusLabel = pyui.widgets.Label('...')
    self.actionFrame.addChild(self.statusLabel)
    
    self.actionFrame.pack()
  
  
  EVENT_ANIMATE = pyui.locals.USEREVENT + 1
  
  def initUI(self):
    """Initialise the UI"""
    pyui.init(self.width, self.height, 'gl', self.fullscreen)
    
    # Initialize glut renderer
    pyui.desktop.getRenderer().setBackMethod(self.render)
    glutReshapeFunc(self.doReshape)
    
    # Initialize keyboard handlers
    self.keyHandler = {
      'a': (self.doAnimation,        'Animate'),
      'c': (self.doCenterView,       'Center view'),
      'd': (self.doSwitchDimensions, 'Switch Dimensions'),
      'h': (self.doToggleHelp,       'Show this help'),
      'i': (self.doToggleInterface,  'Toggle interface'),
      'n': (self.doToggleNeighbours, 'Toggle neighbours'),
      'p': (self.doResetProgress,    'Reset progress'),
      'q': (self.doQuit,             'Quit'),
      'r': (self.initView,           'Reset view'),
      's': (self.doSave,             'Save image'),
    }
    
    pyui.desktop.getDesktop().registerHandler(pyui.locals.KEYDOWN, self.onKeyDown)
    
    self.zooming = False
    pyui.desktop.getDesktop().registerHandler(pyui.locals.KEYUP,   self.onKeyUp)
    
    # Initialize mouse handlers
    pyui.desktop.getDesktop().registerHandler(pyui.locals.LMOUSEBUTTONDOWN, self.onLMouseDown)
    pyui.desktop.getDesktop().registerHandler(pyui.locals.RMOUSEBUTTONDOWN, self.onRMouseDown)
    pyui.desktop.getDesktop().registerHandler(pyui.locals.LMOUSEBUTTONUP,   self.onMouseUp)
    pyui.desktop.getDesktop().registerHandler(pyui.locals.RMOUSEBUTTONUP,   self.onMouseUp)
    
    # Initialize animation handlers
    pyui.desktop.getDesktop().registerHandler(self.EVENT_ANIMATE, self.doAnimation)
    
    # Initialize widgets
    self.initDimFrame()
    self.initActionFrame()
    
    self.interfaceVisible = True
    
    self.initHelpFrame()
  
  def initGL(self):
    """Initialize GL"""
    glClearDepth(1.0)
    glEnable(GL_DEPTH_TEST)
    glClearColor(0.0, 0.0, 0.0, 0.0)
    glShadeModel(GL_SMOOTH)
  
  def initGLData(self):
    """Set up rendering of the vertex_array"""
    # If we also display neighbours, increase point diameter
    glEnable(GL_POINT_SMOOTH)
    glEnable(GL_BLEND)
    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)
    glEnableClientState(GL_VERTEX_ARRAY)
    
    self.updateDisplayedData()
  
  def initGLColors(self, colors):
    """Initialise colored rendering"""
    if colors.shape != (self.n, 3):
      error = "colors argument should be a (%d, 3) array, instead received a %s array."%(self.n, colors.shape)
      raise IndexError, error
    
    if self.k:
      colors = concatenate((colors, ones((self.n * self.k * 2, 3))))
    
    # If specific colors are specified, use them
    glColorPointerf(colors)
    glEnableClientState(GL_COLOR_ARRAY)
  
  
  def doPlay(self, button):
    if self.animating:
      self.playButton.setText('Play', True)
      self.animating = False
    else:
      self.playButton.setText('Pause', True)
      self.animating = True
      
      self.setStatus('Launching animation...')
    
      print self.EVENT_ANIMATE
    
      pyui.desktop.getDesktop().postEvent(self.EVENT_ANIMATE)
  
  def doAnimation(self, event=None):
    """Animate from source to target position by incrementing the progress slider"""
    
    if self.animating:
      # If we haven't received the order not to animate anymore
      self.setStatus("We're animating !!")
    
      self.progressSlider.setValue(self.progressSlider.position + self.speedSlider.position)
      
      self.updateDisplayedData()
    
    if self.progressSlider.position >= self.progressSlider.range:
      self.animating = False
    
    if self.animating:
      # The animation is not finished ; launch next step of the animation
      # after waiting for 1/5 second
      def postAnimateEvent():
        pyui.desktop.getDesktop().postEvent(self.EVENT_ANIMATE)
    
      t = Timer(0.2, postAnimateEvent)
      t.start()
    else:
      self.playButton.setText('Play', True)
      self.setStatus('Done animating.')
  
  def getDispDim(self):
    """
    Returns the currently displayed source and target dimensions, returns
    
      (dispDimsSource, dispDimsTarget)
    
    as in
    
      ([0, 1, 2], [2, 1, 0])
    """
    
    return ([edit.getValue() for edit in self.dimEditSource],
            [edit.getValue() for edit in self.dimEditTarget])
  
  def doSwitchDimensions(self, button=None):
    self.setStatus('Switching dimensions...')
    
    dispDimSource, dispDimTarget = self.getDispDim()
    
    for edit, value in zip(self.dimEditSource, dispDimTarget):
      edit.setValue(value)
    
    for edit, value in zip(self.dimEditTarget, dispDimSource):
      edit.setValue(value)
    
    self.updateDisplayedData()
  
  def doResetProgress(self, button=None):
    self.progressSlider.setValue(0)
  
  def updateDisplayedData(self, value=None):
    """
    Update the OpenGL representation of the data that must be displayed based
    on a linear combination of source and target dimensions given by
    
      source * (1 - progress) + target * progress
    
    We have added minor optimizations for the border cases, when extracting
    data twice from the complete data matrix is not necessary
    (i.e. progress == 0 or 1 - progress == 0)
    """
    
    dispDimsSource = [edit.getValue() for edit in self.dimEditSource]
    dispDimsTarget = [edit.getValue() for edit in self.dimEditTarget]
    
    progress = float(self.progressSlider.position) / self.progressSlider.range
    
    if progress == 0:
      glVertexPointerf(self.data[:, dispDimsSource])
    else:
      if progress == 1:
        glVertexPointerf(self.data[:, dispDimsTarget])
      else:
        glVertexPointerf(
            self.data[:, dispDimsSource] * (1 - progress)
          + self.data[:, dispDimsTarget] * progress)
  
  
  def render(self):
    """Rendering callback"""
    pyui.desktop.getRenderer().clear()
    
    # set up some matrices so that the object spins with the mouse
    glPushMatrix()
    
    gluLookAt(
      0, 0, -1,
      0, 0,  0,
      0, 1,  0)
    glTranslatef(self.transx, self.transy, 0)
    glScalef(self.zoom, self.zoom, self.zoom)
    glRotatef(self.roty, 0, 1, 0)
    glRotatef(self.rotx, 1, 0, 0)
    
    glColor3f(0, 0, 1.0)
    
    if self.k and self.displayNeighbours:
      glPointSize(10)
      
      glDrawArrays(GL_LINES, self.n, self.n * self.k * 2)
    else:
      glPointSize(5)
    
    glDrawArrays(GL_POINTS, 0, self.n)
    
    glBegin(GL_LINES)
    glColor3f(1.0, 0.0, 0.0)
    glVertex4f(0.0, 0.0, 0.0, 1.0)
    glVertex4f(1.0, 0.0, 0.0, 0.0)
    
    glColor3f(0.0, 1.0, 0.0)
    glVertex4f(0.0, 0.0, 0.0, 1.0)
    glVertex4f(0.0, 1.0, 0.0, 0.0)
    
    glColor3f(0.0, 0.0, 1.0)
    glVertex4f(0.0, 0.0, 0.0, 1.0)
    glVertex4f(0.0, 0.0, 1.0, 0.0)
    glEnd()
    
    glPopMatrix()
  
  
  def run(self):
    """Run PyUI main event loop"""
    pyui.run()
    pyui.quit()
  
  
  def onKeyDown(self, event):
    if event.key > 255:
      return 0
    
    if chr(event.key) == 'z':
      if not self.zooming:
        self.zooming = True
        
        self.doMousedown('Zooming', (None, None), self.doZoom)
      
      return 1
    
    try:
      self.keyHandler[chr(event.key)][0]()
      
      return 1
    except Exception, e:
      return 0
  
  def onKeyUp(self, event):
    if event.key > 255:
      return 0
    
    if chr(event.key) == 'z':
      self.zooming = False
      self.onMouseUp(event)
      
      return 1
    
    return 0
  
  def onLMouseDown(self, event):
    self.doMousedown('Rotating', event.pos, self.doRotate)
  
  def onRMouseDown(self, event):
    self.doMousedown('Panning', event.pos, self.doPan)
  
  def doMousedown(self, action, pos, motionHandler):
    self.setStatus('%s...'%action)
    
    self.beginx, self.beginy = pos
    
    pyui.desktop.getDesktop().registerHandler(pyui.locals.MOUSEMOVE, motionHandler)
  
  def onMouseUp(self, event):
    self.setStatus('Waiting...')
    
    pyui.desktop.getDesktop().registerHandler(pyui.locals.MOUSEMOVE, lambda event: None)
  
  def doRotate(self, event):
    x, y = event.pos
    
    self.rotx = self.rotx - (y - self.beginy)
    self.roty = self.roty + (x - self.beginx)
    
    self.beginx, self.beginy = x, y
  
  def doPan(self, event):
    x, y = event.pos
    
    self.transx = self.transx - (x - self.beginx) * self.tScale
    self.transy = self.transy - (y - self.beginy) * self.tScale
    
    self.beginx, self.beginy = x, y
  
  def doZoom(self, event):
    x, y = event.pos
    
    if self.beginx != None:
      self.zoomSlider.setValue(int(self.zoomSlider.position) + (self.beginy - y))
    
    self.beginx, self.beginy = x, y
  
  def changeZoom(self, value):
    self.zoom = max(1, value) / 25.
    
    self.doProjection()
  
  def doReshape(self, width, height):
    """Setup the viewport"""
    
    glViewport(0, 0, width, height)
    
    self.width, self.height = float(width), float(height)
    
    # Setup projection matrix
    self.doProjection()
  
  def doProjection(self):
    """Initialize projection matrix"""
    glMatrixMode(GL_PROJECTION)
    glLoadIdentity()
    
    if self.width > self.height:
      xScale = self.width / self.height
      yScale = 1
      
      self.tScale = 2 * self.bounds / self.height
    else:
      xScale = 1
      yScale = self.height / self.width
      
      self.tScale = 2 * self.bounds / self.width
    
    glOrtho(
      -xScale * self.bounds, xScale * self.bounds,
      -yScale * self.bounds, yScale * self.bounds,
      -self.zoom * self.bounds, 2 * self.zoom * self.bounds)
    
    glMatrixMode(GL_MODELVIEW)
  
  def getViewPort(self):
    return self.width, self.height
  
  
  def setStatus(self, text):
    print text
    self.statusLabel.setText(text)
  
  def doSave(self, button=None):
    import Image
    width, height = self.getViewPort()
    
    filename = 'capture.jpg'
    format   = 'JPEG'
    
    glPixelStorei(GL_PACK_ALIGNMENT, 1)
    data  = glReadPixels(0, 0, width, height, GL_RGB, GL_UNSIGNED_BYTE)
    
    image = Image.fromstring("RGB", (int(width), int(height)), data)
    image = image.transpose(Image.FLIP_TOP_BOTTOM)
    
    image.save(filename, format)
    
    self.setStatus('Saved image to %s'%(os.path.abspath(filename)))
    
    return image
  
  def doQuit(self, button=None):
    self.setStatus('Quitting...')
    
    sys.exit(0)
  
  def doToggleInterface(self):
    self.interfaceVisible = not self.interfaceVisible
    
    self.actionFrame.setShow(self.interfaceVisible)
    self.dimFrame.setShow(self.interfaceVisible)
    
    if not self.interfaceVisible and self.helpVisible:
      self.doToggleHelp()
  
  def doToggleNeighbours(self, button=None):
    self.displayNeighbours = not self.displayNeighbours
  
  def doToggleHelp(self):
    self.helpVisible = not self.helpVisible
    
    self.helpFrame.setShow(self.helpVisible)
  
  def doCenterView(self):
    self.transx = 0
    self.transy = 0
  


if __name__ == '__main__':
  OpenGLDataView(
    array([
      [ 0.0,  0.0,  0.0],
      [10.0,  0.0,  0.0],
      [ 0.0, 10.0,  0.0],
      [ 0.0,  0.0, 10.0],
      [10.0, 10.0,  0.0],
      [ 0.0, 10.0, 10.0],
      [10.0,  0.0, 10.0],
      [10.0, 10.0, 10.0]]),
    knn=3,
    colors=array([[0.0, 0.0, 0.0],
     [1.0, 0.0, 0.0],
     [0.0, 1.0, 0.0],
     [0.0, 0.0, 1.0],
     [1.0, 1.0, 0.0],
     [0.0, 1.0, 1.0],
     [1.0, 0.0, 1.0],
     [1.0, 1.0, 1.0]])
      ).run()
    
  
